package com.company.juc.theadLocal;


import org.apache.log4j.Logger;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;


/**
 * ThreadLocal 实现动态数据源选择
 */
public class DynamicDataSource extends AbstractRoutingDataSource {
    private final static Logger log = Logger.getLogger(DynamicDataSource.class);

    public final static int WRITE_MODE = 0;

    public final static int READ_MODE = 1;

    public final static int RANDOM_MODE = 2;

    private final static Random random = new Random();

    private static final ThreadLocal<Integer> DB_MODE = new ThreadLocal<Integer>();
    /**
     * 数据源的key
     */
    private final List<Object> datasourceKeyList = new ArrayList<Object>();


    /**
     * 获得数据源的key
     */
    @Override
    protected Object determineCurrentLookupKey() {

        // 生成index
        int index = 0;
        int mode = getDBMode();

        if (datasourceKeyList.size() > 1) {
            switch (mode) {
                case WRITE_MODE:
                    index = 0;
                    break;

                case READ_MODE:
                    index = random.nextInt(datasourceKeyList.size() - 1) + 1;
                    break;

                default:
                    index = random.nextInt(datasourceKeyList.size() - 1) + 1;
                    break;
            }
        }

        if (log.isDebugEnabled()) {
            log.debug("determineCurrentLookupKey: " + datasourceKeyList.get(index) + " mode:" + mode);
        }

        return datasourceKeyList.get(index);
    }

    /**
     * 处理数据源的key
     */
    @Override
    protected Object resolveSpecifiedLookupKey(Object lookupKey) {

        datasourceKeyList.add(lookupKey);

        return lookupKey;
    }

    public static void setWriteMode() {
        DB_MODE.set(WRITE_MODE);
    }

    public static void setReadMode() {
        DB_MODE.set(READ_MODE);
    }

    public static int getDBMode() {
        return DB_MODE.get() == null ? RANDOM_MODE : DB_MODE.get();
    }

    public static void clearDBMode() {
        DB_MODE.remove();
    }
}
